import os

import sys
sys.path.append(".")

import glob
import natsort
import cv2

from custom_exp.video_utils import extract_frames_from_video_opencv


def extract_first_frame():
    # video_dir = r"/home/danglingwei/datas/tesseract_test_data/videos"
    video_dir = r"/WORK/PUBLIC/liuyebin_work/lingweidang/datas/tesseract_test_data/videos"
    video_paths = natsort.natsorted(glob.glob(f"{video_dir}/*.mp4"))

    print(f"video_paths: {len(video_paths)}")

    imgs_dir = os.path.join(os.path.dirname(video_dir), "first_frames")
    os.makedirs(imgs_dir, exist_ok=True)


    for vp in video_paths:
        video_name = os.path.splitext(os.path.basename(vp))[0]
        frames = extract_frames_from_video_opencv(vp)
        print(f"video_name: {video_name}, {len(frames)}")

        cv2.imwrite(os.path.join(imgs_dir, f"{video_name}.png"), cv2.cvtColor(frames[0], cv2.COLOR_RGB2BGR))

    print(f"Done!")


if __name__ == "__main__":
    extract_first_frame()
